from RBM import  RBM
import numpy as np
import struct
import cv2
import matplotlib.pyplot as plt
from typing import *


def decode_idx3_ubyte(idx3_ubyte_file):
    """
    解析idx3文件的通用函数
    :param idx3_ubyte_file: idx3文件路径
    :return: 数据集
    """
    # 读取二进制数据
    bin_data = open(idx3_ubyte_file, 'rb').read()

    # 解析文件头信息，依次为魔数、图片数量、每张图片高、每张图片宽
    offset = 0
    fmt_header = '>iiii' #因为数据结构中前4行的数据类型都是32位整型，所以采用i格式，但我们需要读取前4行数据，所以需要4个i。我们后面会看到标签集中，只使用2个ii。
    magic_number, num_images, num_rows, num_cols = struct.unpack_from(fmt_header, bin_data, offset)
    print('魔数:%d, 图片数量: %d张, 图片大小: %d*%d' % (magic_number, num_images, num_rows, num_cols))

    # 解析数据集
    image_size = num_rows * num_cols
    offset += struct.calcsize(fmt_header)  #获得数据在缓存中的指针位置，从前面介绍的数据结构可以看出，读取了前4行之后，指针位置（即偏移位置offset）指向0016。
    print(offset)
    fmt_image = '>' + str(image_size) + 'B'  #图像数据像素值的类型为unsigned char型，对应的format格式为B。这里还有加上图像大小784，是为了读取784个B格式数据，如果没有则只会读取一个值（即一副图像中的一个像素值）
    print(fmt_image,offset,struct.calcsize(fmt_image))
    images = np.empty((num_images, num_rows, num_cols))
    #plt.figure()
    for i in range(num_images):
        if (i + 1) % 10000 == 0:
            print('已解析 %d' % (i + 1) + '张')
            print(offset)
        images[i] = np.array(struct.unpack_from(fmt_image, bin_data, offset)).reshape((num_rows, num_cols))
        #print(images[i])
        offset += struct.calcsize(fmt_image)
        # plt.imshow(images[i],'gray')
        # plt.pause(0.00001)
        # plt.show()
    #plt.show()

    return images



def image2vector(image:np.ndarray)->np.ndarray:
    image[image < 0] = 0
    image[image > 0] = 1
    # plt.imshow(image * 256)
    # plt.show()
    # plt.pause(0)
    return image.flatten()

def vector2image(vector:np.ndarray)->np.ndarray:
    vector = vector.reshape(28,28)
    vector[vector != 1] = 0
    vector[vector == 1] = 256
    return vector


if __name__ == '__main__':
    #实例化模型
    rbm = RBM(784, 100)
    #读取图片数据
    images = decode_idx3_ubyte('./train-images.idx3-ubyte')
    i = 0
    while i < len(images):
        # 构建一批数据
        images_list = images[i:i+100]
        images_batch:List[np.ndarray] = []
        for j in range(len(images_list)):
            image = images_list[j].astype(np.int)
            images_batch.append(image2vector(image))

        #学习一次
        rbm.train(images_batch)
        print("正在处理%s / %s"%(i, len(images)))
        i += 100
        if i >= 10000:
            break


    #模型自动生成一张图片
    for i in range(100):
        v = rbm.generate()
        plt.imshow(vector2image(v), 'gray')
        plt.savefig("./output/%s.png"%(i,))
        print("正在生成 %s.png"%(i,))

